#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
HTTP流量捕获eBPF模块
"""

import os
import time
import logging
import uuid
from datetime import datetime
import socket
import struct
from db.mongo import get_db

# 尝试导入BCC，如果不可用则使用模拟模式
try:
    from bcc import BPF
    BCC_AVAILABLE = True
except ImportError:
    BCC_AVAILABLE = False
    logging.warning("BCC库不可用，将使用模拟模式运行")

logger = logging.getLogger(__name__)

# eBPF程序
bpf_text = """
#include <uapi/linux/ptrace.h>
#include <net/sock.h>
#include <bcc/proto.h>

#define IP_TCP 6
#define ETH_HLEN 14

struct Key {
    u32 src_ip;               // 源IP
    u32 dst_ip;               // 目标IP
    unsigned short src_port;  // 源端口
    unsigned short dst_port;  // 目标端口
};

struct Leaf {
    int timestamp;            // 时间戳
};

// 会话哈希表
BPF_HASH(sessions, struct Key, struct Leaf, 1024);

// 事件数据结构
struct http_event_t {
    u32 src_ip;
    u32 dst_ip;
    u16 src_port;
    u16 dst_port;
    u32 payload_length;
    char payload[4096];  // 最大负载大小
};

// 事件输出通道
BPF_PERF_OUTPUT(events);

// 初始化为0的叶子节点
static struct Leaf zero = {0};

int http_filter(struct __sk_buff *skb) {
    u8 *cursor = 0;

    // 解析以太网头
    struct ethernet_t *ethernet = cursor_advance(cursor, sizeof(*ethernet));

    // 检查是否是IP包
    if (ethernet->type != 0x0800) {
        return 0;
    }

    // 解析IP头
    struct ip_t *ip = cursor_advance(cursor, sizeof(*ip));

    // 检查是否是TCP包
    if (ip->nextp != IP_TCP) {
        return 0;
    }

    // 解析TCP头
    u32 tcp_header_length = 0;
    u32 ip_header_length = 0;
    u32 payload_offset = 0;
    u32 payload_length = 0;

    // 计算IP头长度
    ip_header_length = ip->hlen << 2;

    // 计算TCP头起始位置
    void *tcp_header = (void *)ip + ip_header_length;

    // 解析TCP头
    struct tcp_t *tcp = cursor_advance(cursor, sizeof(*tcp));

    // 计算TCP头长度
    tcp_header_length = tcp->offset << 2;

    // 计算负载起始位置和长度
    payload_offset = ETH_HLEN + ip_header_length + tcp_header_length;
    payload_length = ip->tlen - ip_header_length - tcp_header_length;

    // 创建会话键
    struct Key key = {
        .src_ip = ip->src,
        .dst_ip = ip->dst,
        .src_port = tcp->src_port,
        .dst_port = tcp->dst_port
    };

    // 检查负载长度
    if (payload_length < 7) {
        goto CHECK_SESSION;
    }

    // 加载负载的前7个字节
    unsigned long p[7];
    int i = 0;
    for (i = 0; i < 7; i++) {
        p[i] = load_byte(skb, payload_offset + i);
    }

    // 检查是否是HTTP消息
    // HTTP
    if ((p[0] == 'H') && (p[1] == 'T') && (p[2] == 'T') && (p[3] == 'P')) {
        goto HTTP_MATCH;
    }
    // GET
    if ((p[0] == 'G') && (p[1] == 'E') && (p[2] == 'T')) {
        goto HTTP_MATCH;
    }
    // POST
    if ((p[0] == 'P') && (p[1] == 'O') && (p[2] == 'S') && (p[3] == 'T')) {
        goto HTTP_MATCH;
    }
    // PUT
    if ((p[0] == 'P') && (p[1] == 'U') && (p[2] == 'T')) {
        goto HTTP_MATCH;
    }
    // DELETE
    if ((p[0] == 'D') && (p[1] == 'E') && (p[2] == 'L') && (p[3] == 'E') && (p[4] == 'T') && (p[5] == 'E')) {
        goto HTTP_MATCH;
    }
    // HEAD
    if ((p[0] == 'H') && (p[1] == 'E') && (p[2] == 'A') && (p[3] == 'D')) {
        goto HTTP_MATCH;
    }

    // 检查是否属于已知会话
    CHECK_SESSION:
    struct Leaf *lookup_leaf = sessions.lookup(&key);
    if (lookup_leaf) {
        goto CAPTURE;
    }
    return 0;

    // HTTP匹配，添加到会话
    HTTP_MATCH:
    sessions.lookup_or_try_init(&key, &zero);

    // 捕获数据包
    CAPTURE:
    struct http_event_t event = {};
    event.src_ip = ip->src;
    event.dst_ip = ip->dst;
    event.src_port = tcp->src_port;
    event.dst_port = tcp->dst_port;
    event.payload_length = payload_length > sizeof(event.payload) ? sizeof(event.payload) : payload_length;

    // 复制负载
    for (i = 0; i < event.payload_length; i++) {
        event.payload[i] = load_byte(skb, payload_offset + i);
    }

    // 发送事件
    events.perf_submit(ctx, &event, sizeof(event));

    return 0;
}
"""

class HTTPCapture:
    """HTTP流量捕获类"""

    def __init__(self, interface="eth0"):
        """
        初始化HTTP捕获

        Args:
            interface: 网络接口名称
        """
        self.interface = interface
        self.bpf = None
        self.services = {}  # 服务IP到名称的映射
        self.session_id = None
        self.running = False

    def add_service(self, service_name, service_ip, service_port):
        """
        添加要监控的服务

        Args:
            service_name: 服务名称
            service_ip: 服务IP地址
            service_port: 服务端口
        """
        key = f"{service_ip}:{service_port}"
        self.services[key] = service_name
        logger.info(f"添加监控服务: {service_name} ({service_ip}:{service_port})")

    def start_capture(self, session_id):
        """
        开始捕获HTTP流量

        Args:
            session_id: 会话ID
        """
        if self.running:
            logger.warning("HTTP捕获已经在运行")
            return

        self.session_id = session_id

        try:
            if BCC_AVAILABLE:
                # 加载eBPF程序
                self.bpf = BPF(text=bpf_text)

                # 附加到网络接口
                function = self.bpf.load_func("http_filter", BPF.SOCKET_FILTER)
                BPF.attach_raw_socket(function, self.interface)

                # 设置回调函数
                self.bpf["events"].open_perf_buffer(self._process_event)

                self.running = True
                logger.info(f"开始捕获HTTP流量，会话ID: {session_id}")

                # 启动轮询
                while self.running:
                    self.bpf.perf_buffer_poll(timeout=100)
            else:
                # 模拟模式 - 仅用于测试
                self.running = True
                logger.info(f"使用模拟模式运行，会话ID: {session_id}")

                # 模拟数据生成
                self._generate_mock_data()

                # 保持运行状态，直到被停止
                while self.running:
                    time.sleep(1)

        except Exception as e:
            logger.error(f"HTTP捕获失败: {e}")
            self.running = False

    def _generate_mock_data(self):
        """生成模拟数据用于测试"""
        try:
            # 如果没有服务，则不生成数据
            if not self.services:
                logger.warning("没有服务配置，不生成模拟数据")
                return

            logger.info("生成模拟数据用于测试...")

            # 为每个服务生成一些模拟请求和响应
            for service_key, service_name in self.services.items():
                ip, port = service_key.split(':')

                # 生成一些模拟请求
                for i in range(3):  # 每个服务生成几个请求
                    # 生成请求
                    method = "GET" if i % 2 == 0 else "POST"
                    url = f"/api/{service_name}/endpoint{i}"
                    headers = {
                        "Host": f"{service_name}.example.com",
                        "User-Agent": "Mozilla/5.0",
                        "Content-Type": "application/json"
                    }
                    body = "{}" if method == "GET" else "{\"data\": \"test\"}"

                    # 生成请求ID
                    request_id = str(uuid.uuid4())

                    # 创建请求记录
                    request_data = {
                        'session_id': self.session_id,
                        'service_name': service_name,
                        'request_id': request_id,
                        'timestamp': datetime.now(),
                        'method': method,
                        'url': url,
                        'headers': headers,
                        'body': body,
                        'src_ip': "192.168.1.100",
                        'dst_ip': ip,
                        'src_port': 12345,
                        'dst_port': int(port)
                    }

                    # 生成响应
                    status_code = 200 if i != 1 else 404  # 模拟一些错误
                    response_headers = {
                        "Content-Type": "application/json",
                        "Server": service_name
                    }
                    response_body = "{\"result\": \"success\"}" if status_code == 200 else "{\"error\": \"not found\"}"

                    # 创建响应记录
                    response_data = {
                        'session_id': self.session_id,
                        'service_name': service_name,
                        'request_id': request_id,
                        'timestamp': datetime.now(),
                        'status_code': status_code,
                        'headers': response_headers,
                        'body': response_body,
                        'response_time': 50,  # 模拟响应时间（毫秒）
                        'src_ip': ip,
                        'dst_ip': "192.168.1.100",
                        'src_port': int(port),
                        'dst_port': 12345
                    }

                    # 保存到数据库
                    db = get_db()
                    db.requests.insert_one(request_data)
                    db.responses.insert_one(response_data)

                    logger.info(f"生成模拟数据: {method} {url} -> {service_name}, 状态码: {status_code}")

            logger.info(f"模拟数据生成完成")

        except Exception as e:
            logger.error(f"生成模拟数据失败: {e}")

    def stop_capture(self):
        """停止捕获HTTP流量"""
        self.running = False
        logger.info("停止捕获HTTP流量")

    def _process_event(self, cpu, data, size):
        """
        处理捕获的HTTP事件

        Args:
            cpu: CPU ID
            data: 事件数据
            size: 数据大小
        """
        event = self.bpf["events"].event(data)

        try:
            # 转换IP地址
            src_ip = socket.inet_ntoa(struct.pack("I", event.src_ip))
            dst_ip = socket.inet_ntoa(struct.pack("I", event.dst_ip))
            src_port = event.src_port
            dst_port = event.dst_port

            # 提取负载
            payload = bytes(event.payload[:event.payload_length]).decode('utf-8', errors='replace')

            # 确定服务名称
            src_key = f"{src_ip}:{src_port}"
            dst_key = f"{dst_ip}:{dst_port}"

            service_name = None
            is_request = True

            if src_key in self.services:
                service_name = self.services[src_key]
                is_request = False  # 服务发出的是响应
            elif dst_key in self.services:
                service_name = self.services[dst_key]
                is_request = True  # 发送给服务的是请求

            if not service_name:
                # 未知服务，跳过
                return

            # 解析HTTP内容
            if is_request:
                self._process_request(service_name, payload, src_ip, dst_ip, src_port, dst_port)
            else:
                self._process_response(service_name, payload, src_ip, dst_ip, src_port, dst_port)

        except Exception as e:
            logger.error(f"处理HTTP事件失败: {e}")

    def _process_request(self, service_name, payload, src_ip, dst_ip, src_port, dst_port):
        """处理HTTP请求"""
        try:
            # 解析HTTP请求
            lines = payload.split('\r\n')
            if not lines:
                return

            # 解析请求行
            request_line = lines[0].split(' ')
            if len(request_line) < 3:
                return

            method = request_line[0]
            url = request_line[1]
            version = request_line[2]

            # 解析请求头
            headers = {}
            body = ""
            header_end = False

            for i, line in enumerate(lines[1:]):
                if not line.strip():
                    header_end = True
                    if i + 2 < len(lines):
                        body = '\r\n'.join(lines[i+2:])
                    break

                if not header_end and ':' in line:
                    key, value = line.split(':', 1)
                    headers[key.strip()] = value.strip()

            # 生成请求ID
            request_id = str(uuid.uuid4())

            # 创建请求记录
            request_data = {
                'session_id': self.session_id,
                'service_name': service_name,
                'request_id': request_id,
                'timestamp': datetime.utcnow(),
                'method': method,
                'url': url,
                'headers': headers,
                'body': body,
                'src_ip': src_ip,
                'dst_ip': dst_ip,
                'src_port': src_port,
                'dst_port': dst_port
            }

            # 保存到数据库
            db = get_db()
            db.requests.insert_one(request_data)

            logger.info(f"捕获HTTP请求: {method} {url} -> {service_name}")

            return request_id

        except Exception as e:
            logger.error(f"处理HTTP请求失败: {e}")
            return None

    def _process_response(self, service_name, payload, src_ip, dst_ip, src_port, dst_port):
        """处理HTTP响应"""
        try:
            # 解析HTTP响应
            lines = payload.split('\r\n')
            if not lines:
                return

            # 解析状态行
            status_line = lines[0].split(' ', 2)
            if len(status_line) < 3:
                return

            version = status_line[0]
            status_code = int(status_line[1])
            status_message = status_line[2] if len(status_line) > 2 else ""

            # 解析响应头
            headers = {}
            body = ""
            header_end = False

            for i, line in enumerate(lines[1:]):
                if not line.strip():
                    header_end = True
                    if i + 2 < len(lines):
                        body = '\r\n'.join(lines[i+2:])
                    break

                if not header_end and ':' in line:
                    key, value = line.split(':', 1)
                    headers[key.strip()] = value.strip()

            # 查找对应的请求
            db = get_db()
            request = db.requests.find_one({
                'session_id': self.session_id,
                'service_name': service_name,
                'dst_ip': src_ip,
                'dst_port': src_port
            }, sort=[('timestamp', -1)])

            if not request:
                logger.warning(f"找不到对应的请求: {service_name} {src_ip}:{src_port}")
                return

            request_id = request['request_id']

            # 计算响应时间
            request_time = request['timestamp']
            response_time = datetime.utcnow()
            response_time_ms = int((response_time - request_time).total_seconds() * 1000)

            # 创建响应记录
            response_data = {
                'session_id': self.session_id,
                'service_name': service_name,
                'request_id': request_id,
                'timestamp': response_time,
                'status_code': status_code,
                'headers': headers,
                'body': body,
                'response_time': response_time_ms,
                'src_ip': src_ip,
                'dst_ip': dst_ip,
                'src_port': src_port,
                'dst_port': dst_port
            }

            # 保存到数据库
            db.responses.insert_one(response_data)

            logger.info(f"捕获HTTP响应: {service_name} -> {status_code}")

        except Exception as e:
            logger.error(f"处理HTTP响应失败: {e}")

# 单例实例
http_capture = HTTPCapture()

def start_capture_for_session(session_id):
    """为指定会话启动HTTP捕获"""
    try:
        db = get_db()

        # 获取会话相关的服务
        services = list(db.services.find({'session_id': session_id}))

        if not services:
            logger.warning(f"会话 {session_id} 没有关联的服务")
            return False

        # 添加服务到捕获器
        for service in services:
            http_capture.add_service(
                service['service_name'],
                service['service_ip'],
                service['service_port']
            )

        # 如果在模拟模式下，直接生成模拟数据
        if not BCC_AVAILABLE:
            http_capture.session_id = session_id
            http_capture._generate_mock_data()
            logger.info(f"在模拟模式下生成了测试数据，会话ID: {session_id}")
            return True

        # 启动捕获
        # 在新线程中启动，避免阻塞API
        import threading
        capture_thread = threading.Thread(
            target=http_capture.start_capture,
            args=(session_id,),
            daemon=True
        )
        capture_thread.start()

        logger.info(f"在后台启动HTTP捕获线程，会话ID: {session_id}")
        return True

    except Exception as e:
        logger.error(f"启动HTTP捕获失败: {e}")
        return False

def stop_capture():
    """停止HTTP捕获"""
    http_capture.stop_capture()
    return True
